
import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from pylab import array
import math
from prepocess.utils import get_file


def plot_batch_img(img_arry_list,img_label_list):
    border =  math.ceil(math.sqrt(len(img_arry_list)))
    # plt.suptitle('plot batch')
    for x in range(1,border*border):
        if x > len(img_arry_list):
            break
        plt.subplot(border,border,x) , plt.title(str(img_label_list[x-1]))
        plt.imshow(img_arry_list[x-1]), plt.axis('off')
        plt.tight_layout()
    plt.show()


def convert_ids_to_class_names_onehot(ids,ids_to_class_names):
    return [ ids_to_class_names[np.argmax(id)] for id in ids]


def convert_ids_to_class_names(ids,ids_to_class_names):
    return [ ids_to_class_names[id] for id in ids]


def convert_img_path_to_img_array_list(img_path_list):
    return [ array(Image.open(path)) for path in img_path_list]



def case1():
    img = Image.open("test.PNG")
    gray = img.convert('L')
    r, g, b, k = img.split()
    img_merged = Image.merge('RGB', (r, g, b))

    plt.figure(figsize=(10, 5))  # 设置窗口大小
    plt.suptitle('Multi_Image')  # 图片名称
    plt.subplot(2, 3, 1), plt.title('image')
    plt.imshow(img), plt.axis('off')
    plt.subplot(2, 3, 2), plt.title('gray')
    plt.imshow(gray, cmap='gray'), plt.axis('off')  # 这里显示灰度图要加cmap
    plt.subplot(2, 3, 3), plt.title('img_merged')
    plt.imshow(img_merged), plt.axis('off')
    plt.subplot(2, 3, 4), plt.title('r')
    plt.imshow(r, cmap='gray'), plt.axis('off')
    plt.subplot(2, 3, 5), plt.title('g')
    plt.imshow(g, cmap='gray'), plt.axis('off')
    plt.subplot(2, 3, 6), plt.title('b')
    plt.imshow(b, cmap='gray'), plt.axis('off')
    plt.show()


def case2():
    tmp = Image.open('test.PNG')
    im = array(tmp)
    im_list = [im, im, im, im, im, im, im]
    label_list = ['car', 'bike', 'person', 'car', 'bike', 'person', 'plane']
    plot_batch_img(im_list, label_list)


def case3():
    train_dir = '../data_set_lit/'  # 训练样本的读入路径
    train, train_label, val, val_label, class_names_to_ids, ids_to_class_names = get_file(train_dir)

    plot_batch_img(convert_img_path_to_img_array_list(train),
                   convert_ids_to_class_names(train_label, ids_to_class_names))


if __name__ == '__main__':
    case1()
    case2()
    case3()
